"""
1D Diffusion Equation solver using locELM with block time marching.

This script reproduces the experiments from the paper for the 1D diffusion equation:
    du/dt - nu * d^2u/dx^2 = f(x,t)

Domain: x in [0, 5], t in [0, t_f]
Parameters: nu = 0.01, t_f = 10 or 1
Exact solution: u(x,t) = [2*cos(pi*x + pi/5) + 1.5*cos(2*pi*x - 3*pi/5)] *
                          [2*cos(pi*t + pi/5) + 1.5*cos(2*pi*t - 3*pi/5)]
"""

import numpy as np
import tensorflow as tf
import matplotlib.pyplot as plt
import time
from scipy.linalg import lstsq
import sys
import os

# Add parent directory to path
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from locELM.core.domain import Domain
from locELM.core.networks import LocalELMNetwork, MultiSubdomainNetwork
from locELM.core.continuity import ContinuityConditions
from locELM.utils.helpers import create_evaluation_function


class Diffusion1DSolver:
    """Solver for 1D diffusion equation using locELM with block time marching."""

    def __init__(self, nu=0.01, a=0.0, b=5.0, t_f=10.0, seed=22):
        """Initialize solver parameters."""
        self.nu = nu
        self.a = a
        self.b = b
        self.t_f = t_f
        self.seed = seed

        # Set random seed
        tf.random.set_seed(seed)
        np.random.seed(seed)

    def exact_solution(self, x, t):
        """Compute exact solution."""
        spatial = (2 * np.cos(np.pi * x + np.pi/5) +
                   1.5 * np.cos(2*np.pi * x - 3*np.pi/5))
        temporal = (2 * np.cos(np.pi * t + np.pi/5) +
                    1.5 * np.cos(2*np.pi * t - 3*np.pi/5))
        return spatial * temporal

    def source_term(self, x, t):
        """Compute source term f(x,t) such that exact solution satisfies PDE."""
        # du/dt
        spatial = (2 * np.cos(np.pi * x + np.pi/5) +
                   1.5 * np.cos(2*np.pi * x - 3*np.pi/5))
        dt_temporal = (-2*np.pi * np.sin(np.pi * t + np.pi/5) -
                       3*np.pi * np.sin(2*np.pi * t - 3*np.pi/5))
        dudt = spatial * dt_temporal

        # d^2u/dx^2
        dx2_spatial = (-2*np.pi**2 * np.cos(np.pi * x + np.pi/5) -
                       6*np.pi**2 * np.cos(2*np.pi * x - 3*np.pi/5))
        temporal = (2 * np.cos(np.pi * t + np.pi/5) +
                    1.5 * np.cos(2*np.pi * t - 3*np.pi/5))
        d2udx2 = dx2_spatial * temporal

        # f = du/dt - nu * d^2u/dx^2
        return dudt - self.nu * d2udx2

    def boundary_condition_left(self, t):
        """Boundary condition at x = a."""
        return self.exact_solution(self.a, t)

    def boundary_condition_right(self, t):
        """Boundary condition at x = b."""
        return self.exact_solution(self.b, t)

    def initial_condition(self, x):
        """Initial condition at t = 0."""
        return self.exact_solution(x, 0.0)

    def solve_time_block(self, t_start, t_end, initial_data,
                         n_x=5, n_t=1, q_x=30, q_t=30, M=300, R_m=1.0):
        """
        Solve one time block.

        Parameters
        ----------
        t_start : float
            Start time of block
        t_end : float
            End time of block
        initial_data : callable
            Function u(x) at t=t_start
        n_x : int
            Number of subdomains in x direction
        n_t : int
            Number of subdomains in t direction
        q_x : int
            Number of collocation points in x per subdomain
        q_t : int
            Number of collocation points in t per subdomain
        M : int
            Number of training parameters (hidden layer width)
        R_m : float
            Maximum magnitude of random coefficients

        Returns
        -------
        weights : np.ndarray
            Trained weights
        networks : MultiSubdomainNetwork
            Trained networks
        max_error : float
            Maximum error in block
        rms_error : float
            RMS error in block
        train_time : float
            Training time
        """
        start_time = time.time()

        # Create domain (2D: x and t)
        bounds = np.array([[self.a, self.b], [t_start, t_end]])
        n_subdomains = [n_x, n_t]
        domain = Domain(bounds, n_subdomains)

        # Create networks
        networks = MultiSubdomainNetwork(
            n_subdomains=domain.n_total,
            input_dim=2,
            output_dim=1,
            hidden_layers=[M],
            activation='tanh',
            r_max=R_m,
            seed=self.seed
        )

        # Add networks for each subdomain
        for idx in range(domain.n_total):
            region = domain.get_subdomain_region(idx)
            networks.add_subdomain_network(region)

        # Generate collocation points
        collocation_points = {}
        for idx in range(domain.n_total):
            region = domain.get_subdomain_region(idx)
            # Uniform grid
            x_points = np.linspace(region[0][0], region[0][1], q_x)
            t_points = np.linspace(region[1][0], region[1][1], q_t)
            X, T = np.meshgrid(x_points, t_points)
            points = np.column_stack([X.ravel(), T.ravel()])
            collocation_points[idx] = points

        # Build system of equations
        equations = []
        rhs = []

        # PDE equations
        for subdomain_idx in range(domain.n_total):
            points = collocation_points[subdomain_idx]
            x = points[:, 0]
            t = points[:, 1]
            network = networks.get_network(subdomain_idx)

            # Compute diffusion operator on last hidden layer
            LV = self._diffusion_operator(network, points)
            f = self.source_term(x, t)

            for i in range(len(points)):
                eq_row = np.zeros(networks.n_total_params)
                idx_start = subdomain_idx * M
                eq_row[idx_start:idx_start + M] = LV[i, :]
                equations.append(eq_row)
                rhs.append(f[i])

        # Boundary conditions (spatial boundaries)
        for subdomain_idx in range(domain.n_total):
            points = collocation_points[subdomain_idx]
            network = networks.get_network(subdomain_idx)

            # Left boundary (x = a)
            mask_left = np.abs(points[:, 0] - self.a) < 1e-10
            if np.any(mask_left):
                boundary_points = points[mask_left]
                V = network.compute_last_hidden_output(boundary_points)
                g = self.boundary_condition_left(boundary_points[:, 1])

                for i in range(len(boundary_points)):
                    eq_row = np.zeros(networks.n_total_params)
                    idx_start = subdomain_idx * M
                    eq_row[idx_start:idx_start + M] = V[i, :]
                    equations.append(eq_row)
                    rhs.append(g[i])

            # Right boundary (x = b)
            mask_right = np.abs(points[:, 0] - self.b) < 1e-10
            if np.any(mask_right):
                boundary_points = points[mask_right]
                V = network.compute_last_hidden_output(boundary_points)
                g = self.boundary_condition_right(boundary_points[:, 1])

                for i in range(len(boundary_points)):
                    eq_row = np.zeros(networks.n_total_params)
                    idx_start = subdomain_idx * M
                    eq_row[idx_start:idx_start + M] = V[i, :]
                    equations.append(eq_row)
                    rhs.append(g[i])

        # Initial condition (t = t_start)
        for subdomain_idx in range(domain.n_total):
            points = collocation_points[subdomain_idx]
            network = networks.get_network(subdomain_idx)

            mask_initial = np.abs(points[:, 1] - t_start) < 1e-10
            if np.any(mask_initial):
                boundary_points = points[mask_initial]
                V = network.compute_last_hidden_output(boundary_points)
                h = initial_data(boundary_points[:, 0])

                for i in range(len(boundary_points)):
                    eq_row = np.zeros(networks.n_total_params)
                    idx_start = subdomain_idx * M
                    eq_row[idx_start:idx_start + M] = V[i, :]
                    equations.append(eq_row)
                    rhs.append(h[i])

        # Continuity conditions
        continuity = ContinuityConditions(domain, networks, continuity_order=[1, 0])  # C^1 in x, C^0 in t
        A_cont, b_cont = continuity.build_continuity_equations(collocation_points, output_components=1)

        # Stack all equations
        A = np.vstack([np.array(equations), A_cont])
        b = np.hstack([np.array(rhs), b_cont])

        # Solve least squares
        weights, residual, rank, s = lstsq(A, b, lapack_driver='gelsy')
        networks.set_all_weights(weights)

        train_time = time.time() - start_time

        # Compute errors
        eval_func = create_evaluation_function(networks, domain, output_dim=1)

        errors = []
        for subdomain_idx in range(domain.n_total):
            region = domain.get_subdomain_region(subdomain_idx)
            x_test = np.linspace(region[0][0], region[0][1], 50)
            t_test = np.linspace(region[1][0], region[1][1], 50)
            X_test, T_test = np.meshgrid(x_test, t_test)
            points_test = np.column_stack([X_test.ravel(), T_test.ravel()])

            u_exact = self.exact_solution(points_test[:, 0], points_test[:, 1])
            u_pred = eval_func(points_test)

            errors.extend(np.abs(u_exact - u_pred.flatten()))

        max_error = np.max(errors)
        rms_error = np.sqrt(np.mean(np.array(errors)**2))

        return weights, networks, domain, max_error, rms_error, train_time

    def _diffusion_operator(self, network, points):
        """
        Apply diffusion operator: du/dt - nu * d^2u/dx^2

        Returns operator applied to last hidden layer outputs.
        """
        x = tf.constant(points, dtype=tf.float32)
        M = network.hidden_layers[-1]
        n_points = points.shape[0]

        LV = np.zeros((n_points, M))

        for j in range(M):
            # Compute du/dt
            with tf.GradientTape() as tape:
                tape.watch(x)
                V = network.last_hidden_model(x, training=False)
                V_j = V[:, j]
            dV_dx_and_dt = tape.gradient(V_j, x)
            dV_dt = dV_dx_and_dt[:, 1]

            # Compute d^2u/dx^2
            with tf.GradientTape(persistent=True) as tape2:
                tape2.watch(x)
                with tf.GradientTape() as tape1:
                    tape1.watch(x)
                    V = network.last_hidden_model(x, training=False)
                    V_j = V[:, j]
                dV_dx = tape1.gradient(V_j, x)[:, 0]
            d2V_dx2 = tape2.gradient(dV_dx, x)[:, 0]

            # Diffusion operator: dV/dt - nu * d^2V/dx^2
            LV[:, j] = dV_dt.numpy() - self.nu * d2V_dx2.numpy()

        return LV

    def solve_with_time_blocks(self, n_blocks=10, n_x=5, n_t=1,
                               q_x=30, q_t=30, M=300, R_m=1.0):
        """
        Solve using block time marching.

        Parameters
        ----------
        n_blocks : int
            Number of time blocks
        n_x : int
            Number of subdomains in x per block
        n_t : int
            Number of subdomains in t per block
        q_x : int
            Collocation points in x per subdomain
        q_t : int
            Collocation points in t per subdomain
        M : int
            Training parameters per subdomain
        R_m : float
            Random coefficient magnitude

        Returns
        -------
        dict
            Results including errors, times, and solution data
        """
        dt_block = self.t_f / n_blocks

        all_max_errors = []
        all_rms_errors = []
        all_train_times = []
        all_networks = []

        # Initial condition for first block
        initial_data = self.initial_condition

        for block_idx in range(n_blocks):
            t_start = block_idx * dt_block
            t_end = (block_idx + 1) * dt_block

            print(f"Solving time block {block_idx + 1}/{n_blocks}: t=[{t_start:.2f}, {t_end:.2f}]")

            weights, networks, domain, max_err, rms_err, train_time = self.solve_time_block(
                t_start, t_end, initial_data, n_x, n_t, q_x, q_t, M, R_m
            )

            all_max_errors.append(max_err)
            all_rms_errors.append(rms_err)
            all_train_times.append(train_time)
            all_networks.append(networks)

            print(f"  Max error: {max_err:.3e}, RMS error: {rms_err:.3e}, Time: {train_time:.2f}s")

            # Update initial condition for next block (use solution at t_end)
            # Create evaluation function for this block
            block_eval_func = create_evaluation_function(networks, domain, output_dim=1)

            def next_initial_data(x):
                # Evaluate at t=t_end
                points = np.column_stack([x, np.full_like(x, t_end)])
                return block_eval_func(points)

            initial_data = next_initial_data

        overall_max_error = np.max(all_max_errors)
        overall_rms_error = np.sqrt(np.mean(np.array(all_rms_errors)**2))
        total_train_time = np.sum(all_train_times)

        results = {
            'n_blocks': n_blocks,
            'n_x': n_x,
            'n_t': n_t,
            'q_x': q_x,
            'q_t': q_t,
            'M': M,
            'R_m': R_m,
            'max_errors': all_max_errors,
            'rms_errors': all_rms_errors,
            'train_times': all_train_times,
            'overall_max_error': overall_max_error,
            'overall_rms_error': overall_rms_error,
            'total_train_time': total_train_time,
            'networks': all_networks
        }

        return results


if __name__ == '__main__':
    # Create solver
    solver = Diffusion1DSolver(nu=0.01, a=0.0, b=5.0, t_f=10.0, seed=22)

    # Test case from Figure 1 in paper: 10 time blocks, 5 subdomains, 30x30 collocation, 300 training params
    print("=" * 60)
    print("1D Diffusion Equation - Reproducing Figure 1")
    print("=" * 60)
    results = solver.solve_with_time_blocks(
        n_blocks=10,
        n_x=5,
        n_t=1,
        q_x=30,
        q_t=30,
        M=300,
        R_m=1.0
    )

    print("\n" + "=" * 60)
    print("FINAL RESULTS")
    print("=" * 60)
    print(f"Overall Maximum Error: {results['overall_max_error']:.3e}")
    print(f"Overall RMS Error: {results['overall_rms_error']:.3e}")
    print(f"Total Training Time: {results['total_train_time']:.2f} seconds")
    print("=" * 60)
